#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Main script for RINN (Rank Inspired Neural Network) implementation.
Orchestrates PIELM, RINN, and RINN-es methods with comparison capabilities.

Based on Paper 2506.17654 - "Rank Inspired Neural Network for solving linear partial differential equations"
"""

import torch
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for headless plotting
import argparse
import sys
import os
import time
from pathlib import Path

# Set working directory to current script directory and add to Python path
script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)
sys.path.insert(0, script_dir)

# Import RINN modules
from pielm.solver import create_pielm_solver, laplacian_operator, helmholtz_operator
from rinn.pretrainer import RINNPretrainer
from rinn.rinn_es import create_rinn_es_solver, RINNComparison
from experiments.poisson_2d_high_freq import PoissonHighFreq2D, run_convergence_study, demonstrate_solution_quality
from common.utils import set_random_seed, generate_collocation_points, generate_boundary_points
from common.metrics import evaluate_model_comprehensive, print_metrics
from common.plotting import plot_comparison, plot_training_history


def create_simple_poisson_problem():
    """Create a simple 2D Poisson problem for testing."""
    
    class SimplePoisson:
        def __init__(self):
            self.domain = {'x_min': 0.0, 'x_max': 1.0, 'y_min': 0.0, 'y_max': 1.0}
        
        def analytical_solution(self, points):
            """u(x,y) = x(1-x)y(1-y)"""
            x = points[:, 0:1]
            y = points[:, 1:2]
            return x * (1 - x) * y * (1 - y)
        
        def source_function(self, points):
            """f = -nabla^2 u = 2x(1-x) + 2y(1-y)"""
            x = points[:, 0:1]
            y = points[:, 1:2]
            return 2 * x * (1 - x) + 2 * y * (1 - y)
        
        def boundary_values(self, points):
            """Homogeneous Dirichlet BC (u=0 on boundary)"""
            return torch.zeros(points.shape[0], 1)
    
    return SimplePoisson()


def run_pielm_demo(hidden_dim=100, seed=42):
    """Demonstrate PIELM method."""
    print(f"\n{'='*60}")
    print("PIELM DEMONSTRATION")
    print('='*60)
    
    set_random_seed(seed)
    problem = create_simple_poisson_problem()
    
    # Generate training data
    interior_points = generate_collocation_points(
        domain=problem.domain, n_points=1000, method='random'
    )
    boundary_points = generate_boundary_points(
        domain=problem.domain, n_boundary=200
    )
    
    source_values = problem.source_function(interior_points)
    boundary_values = problem.boundary_values(boundary_points)
    
    print(f"Training PIELM with {hidden_dim} hidden neurons...")
    print(f"Interior points: {len(interior_points)}, Boundary points: {len(boundary_points)}")
    
    # Create and train PIELM
    start_time = time.time()
    pielm_solver = create_pielm_solver(
        input_dim=2,
        hidden_dim=hidden_dim,
        activation='tanh',
        seed=seed
    )
    
    training_info = pielm_solver.train(
        collocation_points=interior_points,
        boundary_points=boundary_points,
        pde_operator=laplacian_operator,
        pde_rhs=source_values,
        boundary_values=boundary_values
    )
    training_time = time.time() - start_time
    
    print(f"Training completed in {training_time:.2f}s")
    print(f"System condition number: {training_info['condition_number']:.6e}")
    print(f"Relative residual: {training_info['relative_residual']:.6e}")
    
    # Evaluate on test data
    test_points = generate_collocation_points(
        domain=problem.domain, n_points=2500, method='uniform'
    )
    true_solution = problem.analytical_solution(test_points)
    
    metrics = evaluate_model_comprehensive(
        network=pielm_solver.network,
        test_points=test_points,
        true_solution=true_solution,
        pde_operator=laplacian_operator
    )
    
    print_metrics(metrics, "PIELM Evaluation")
    
    return {
        'solver': pielm_solver,
        'training_info': training_info,
        'metrics': metrics,
        'test_points': test_points,
        'true_solution': true_solution
    }


def run_rinn_pretraining_demo(hidden_dim=100, seed=42):
    """Demonstrate RINN pretraining (Stage 1)."""
    print(f"\n{'='*60}")
    print("RINN PRETRAINING DEMONSTRATION")
    print('='*60)
    
    set_random_seed(seed)
    
    from pielm.solver import PIELMNetwork
    
    # Create network
    network = PIELMNetwork(input_dim=2, hidden_dim=hidden_dim, activation='tanh')
    
    # Generate collocation points for pretraining
    domain = {'x_min': 0.0, 'x_max': 1.0, 'y_min': 0.0, 'y_max': 1.0}
    points = generate_collocation_points(domain, n_points=2000, method='random')
    
    print(f"Pretraining RINN with {hidden_dim} hidden neurons...")
    print(f"Collocation points: {len(points)}")
    
    # Create pretrainer
    pretrainer = RINNPretrainer(
        network=network,
        epsilon=1e-2,
        learning_rate=1e-3,
        optimizer_type='adam'
    )
    
    # Evaluate orthogonality before training
    print("\nOrthogonality before pretraining:")
    metrics_before = pretrainer.evaluate_orthogonality(points)
    for name, value in metrics_before.items():
        print(f"  {name}: {value:.6e}")
    
    # Train pretrainer
    start_time = time.time()
    history = pretrainer.train(
        points=points,
        max_epochs=1000,
        print_every=200,
        tolerance=1e-6,
        patience=100
    )
    training_time = time.time() - start_time
    
    print(f"Pretraining completed in {training_time:.2f}s")
    
    # Evaluate orthogonality after training
    print("\nOrthogonality after pretraining:")
    metrics_after = pretrainer.evaluate_orthogonality(points)
    for name, value in metrics_after.items():
        print(f"  {name}: {value:.6e}")
    
    print(f"\nImprovement in off-diagonal norm: "
          f"{metrics_before['off_diagonal_norm']:.6e} -> {metrics_after['off_diagonal_norm']:.6e} "
          f"({metrics_before['off_diagonal_norm']/metrics_after['off_diagonal_norm']:.1f}x reduction)")
    
    return {
        'pretrainer': pretrainer,
        'network': network,
        'history': history,
        'metrics_before': metrics_before,
        'metrics_after': metrics_after
    }


def run_rinn_es_demo(hidden_dim=100, seed=42):
    """Demonstrate RINN-es method."""
    print(f"\n{'='*60}")
    print("RINN-ES DEMONSTRATION")
    print('='*60)
    
    set_random_seed(seed)
    problem = create_simple_poisson_problem()
    
    # Generate training data
    interior_points = generate_collocation_points(
        domain=problem.domain, n_points=1000, method='random'
    )
    boundary_points = generate_boundary_points(
        domain=problem.domain, n_boundary=200
    )
    
    source_values = problem.source_function(interior_points)
    boundary_values = problem.boundary_values(boundary_points)
    
    print(f"Training RINN-es with {hidden_dim} hidden neurons...")
    print(f"Interior points: {len(interior_points)}, Boundary points: {len(boundary_points)}")
    
    # Create and train RINN-es
    start_time = time.time()
    rinn_es = create_rinn_es_solver(
        input_dim=2,
        hidden_dim=hidden_dim,
        activation='tanh',
        epsilon=1e-2,
        learning_rate=1e-3,
        patience=30,
        seed=seed
    )
    
    training_results = rinn_es.train(
        collocation_points=interior_points,
        boundary_points=boundary_points,
        pde_operator=laplacian_operator,
        pde_rhs=source_values,
        boundary_values=boundary_values,
        max_epochs=200,
        print_every=20
    )
    training_time = time.time() - start_time
    
    print(f"Training completed in {training_time:.2f}s")
    print(f"Best PDE loss: {training_results['best_pde_loss']:.6e} at epoch {training_results['best_epoch']}")
    print(f"Early stopping: {'Yes' if training_results['stopped_early'] else 'No'}")
    
    # Evaluate on test data
    test_points = generate_collocation_points(
        domain=problem.domain, n_points=2500, method='uniform'
    )
    true_solution = problem.analytical_solution(test_points)
    
    metrics = rinn_es.evaluate_final_state(test_points, true_solution)
    print_metrics(metrics, "RINN-es Evaluation")
    
    return {
        'solver': rinn_es,
        'training_results': training_results,
        'metrics': metrics,
        'test_points': test_points,
        'true_solution': true_solution
    }


def run_method_comparison(hidden_dims=[50, 100, 200], n_trials=3):
    """Compare PIELM and RINN-es methods."""
    print(f"\n{'='*60}")
    print("METHOD COMPARISON")
    print('='*60)
    
    problem = create_simple_poisson_problem()
    
    # Prepare problem configuration
    interior_points = generate_collocation_points(
        domain=problem.domain, n_points=1000, method='random'
    )
    boundary_points = generate_boundary_points(
        domain=problem.domain, n_boundary=200
    )
    source_values = problem.source_function(interior_points)
    boundary_values = problem.boundary_values(boundary_points)
    
    test_points = generate_collocation_points(
        domain=problem.domain, n_points=2500, method='uniform'
    )
    true_solution = problem.analytical_solution(test_points)
    
    problem_config = {
        'collocation_points': interior_points,
        'boundary_points': boundary_points,
        'pde_operator': laplacian_operator,
        'pde_rhs': source_values,
        'boundary_values': boundary_values,
        'test_points': test_points,
        'true_solution': true_solution
    }
    
    network_configs = [{'hidden_dim': hd} for hd in hidden_dims]
    
    # Run comparison
    results = RINNComparison.run_comparison(
        problem_config=problem_config,
        network_configs=network_configs,
        n_trials=n_trials
    )
    
    # Print summary
    RINNComparison.summarize_results(results)
    
    return results


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='RINN (Rank Inspired Neural Network) Implementation',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python main.py --demo pielm                    # Run PIELM demo
  python main.py --demo rinn-pretraining         # Run RINN pretraining demo
  python main.py --demo rinn-es                  # Run RINN-es demo
  python main.py --demo all                      # Run all demos
  python main.py --comparison                    # Compare methods
  python main.py --experiment poisson-2d         # Run Poisson experiment
  python main.py --experiment convergence        # Run convergence study
        """
    )
    
    parser.add_argument('--demo', type=str, choices=['pielm', 'rinn-pretraining', 'rinn-es', 'all'],
                        help='Run demonstration for specific method')
    parser.add_argument('--comparison', action='store_true',
                        help='Run method comparison')
    parser.add_argument('--experiment', type=str, choices=['poisson-2d', 'convergence'],
                        help='Run specific experiment')
    parser.add_argument('--hidden-dim', type=int, default=100,
                        help='Number of hidden neurons (default: 100)')
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed (default: 42)')
    parser.add_argument('--frequency', type=int, default=4,
                        help='Frequency parameter for Poisson experiment (default: 4)')
    parser.add_argument('--n-trials', type=int, default=3,
                        help='Number of trials for comparison/convergence (default: 3)')
    
    args = parser.parse_args()
    
    # Print header
    print("="*80)
    print("RINN: Rank Inspired Neural Network Implementation")
    print("Paper 2506.17654 Reproduction")
    print("="*80)
    
    if args.demo:
        if args.demo in ['pielm', 'all']:
            run_pielm_demo(hidden_dim=args.hidden_dim, seed=args.seed)
        
        if args.demo in ['rinn-pretraining', 'all']:
            run_rinn_pretraining_demo(hidden_dim=args.hidden_dim, seed=args.seed)
        
        if args.demo in ['rinn-es', 'all']:
            run_rinn_es_demo(hidden_dim=args.hidden_dim, seed=args.seed)
    
    elif args.comparison:
        hidden_dims = [50, 100, 200] if args.hidden_dim == 100 else [args.hidden_dim]
        run_method_comparison(hidden_dims=hidden_dims, n_trials=args.n_trials)
    
    elif args.experiment == 'poisson-2d':
        demonstrate_solution_quality(frequency=args.frequency, hidden_dim=args.hidden_dim)
    
    elif args.experiment == 'convergence':
        hidden_dims = [50, 100, 200, 400]
        run_convergence_study(
            frequency=args.frequency,
            hidden_dims=hidden_dims,
            n_trials=args.n_trials
        )
    
    else:
        parser.print_help()
        print("\nNo action specified. Use --demo, --comparison, or --experiment.")
    
    print(f"\n{'='*80}")
    print("Execution completed!")
    print("="*80)


if __name__ == "__main__":
    main()